from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel
import numpy as np
import uuid
import base64
import io
from PIL import Image

from config import get_config
from logger import setup_logger
from model import sam_model

config = get_config()
logger = setup_logger('sam_api')

app = FastAPI(
    title="SAM Embedding Service",
    description="Segment Anything Model (SAM) 图像嵌入服务，支持GPU加速",
    version="1.0.0"
)

# 添加CORS支持
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境应限制具体域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法
    allow_headers=["*"],  # 允许所有HTTP头
)

class EmbeddingResponse(BaseModel):
    """embedding响应模型"""
    embedding: list  # embedding数据
    shape: list      # embedding形状
    image_id: str    # 图像ID
    model_type: str  # 模型类型

class ImageRequest(BaseModel):
    """图像请求模型"""
    image_base64: str  # base64编码的图像数据

@app.get("/")
def root():
    """根路径，返回服务信息"""
    return {
        "service": "SAM Embedding Service",
        "version": "1.0.0",
        "status": "running",
        "device": config.DEVICE,
        "model_type": config.MODEL_TYPE
    }

@app.get("/health")
def health_check():
    """健康检查接口"""
    try:
        # 检查模型是否加载成功
        if not sam_model.model_loaded:
            loaded = sam_model.load_model()
            if not loaded:
                return JSONResponse(
                    status_code=503,
                    content={"status": "error", "message": "模型加载失败"}
                )
        return {"status": "ok", "device": config.DEVICE, "model_type": config.MODEL_TYPE}
    except Exception as e:
        logger.error(f"健康检查失败: {e}")
        return JSONResponse(
            status_code=503,
            content={"status": "error", "message": str(e)}
        )

# 保留原始的/embedding接口，现在作为文件上传接口
@app.post("/embedding", response_model=EmbeddingResponse)
async def generate_embedding(
    file: UploadFile = File(..., description="要处理的图像文件")
):
    """上传图像文件，生成embedding"""
    try:
        # 检查文件类型
        if not file.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="上传的文件必须是图像")
        
        # 读取文件内容
        image_bytes = await file.read()
        
        # 生成唯一的图像ID
        image_id = str(uuid.uuid4())
        
        logger.info(f"接收到图像请求，文件: {file.filename}, 大小: {len(image_bytes)} bytes, ID: {image_id}")
        
        # 生成embedding
        embedding = sam_model.generate_embedding(image_bytes, image_id)
        
        # 将numpy数组转换为列表
        embedding_list = embedding.tolist()
        
        logger.info(f"成功生成embedding，形状: {embedding.shape}, ID: {image_id}")
        
        return EmbeddingResponse(
            embedding=embedding_list,
            shape=list(embedding.shape),
            image_id=image_id,
            model_type=config.MODEL_TYPE
        )
    except Exception as e:
        logger.error(f"生成embedding失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# /embedding/base64接口，返回.npy格式二进制流
@app.post("/embedding/base64")
async def generate_embedding_base64(
    image_base64: str = Form(..., description="base64编码的图像数据")
):
    """通过base64字符串生成图像embedding，返回.npy格式的二进制流
    该函数接收一个Base64编码的字符串，将其解码为图像，并使用预训练的模型生成该图像的嵌入向量。
    生成的嵌入向量以二进制流的形式返回，文件格式为.npy，适用于后续的深度学习任务。
    
    参数:
    image_base64 (str): Base64编码的图像字符串。
    
    返回:
    Response: 包含图像嵌入向量的二进制流响应，文件格式为.npy。
    响应头中包含图像ID、模型类型和embedding形状等信息。
    """
    try:
        logger.info(f"接收到base64图像请求，长度: {len(image_base64)}字符")
        
        # 生成唯一的图像ID
        image_id = str(uuid.uuid4())
        
        # 解码base64字符串为字节
        try:
            image_bytes = base64.b64decode(image_base64)
            logger.info(f"base64解码成功，图像大小: {len(image_bytes)} bytes")
        except Exception as e:
            logger.error(f"base64解码失败: {e}")
            raise HTTPException(status_code=400, detail="无效的base64编码")
        
        # 验证是否为图像
        try:
            Image.open(io.BytesIO(image_bytes))
        except Exception as e:
            logger.error(f"图像验证失败: {e}")
            raise HTTPException(status_code=400, detail="解码后的内容不是有效的图像")
        
        # 生成embedding
        embedding = sam_model.generate_embedding(image_bytes, image_id)
        logger.info(f"成功生成embedding，形状: {embedding.shape}, ID: {image_id},dtype:{embedding.dtype}")
        # 检查embedding是否已经是numpy数组
        if hasattr(embedding, 'cpu'):
            # 如果是PyTorch张量，则转换为numpy数组
            image_embedding = embedding.cpu().numpy()
        else:
            # 如果已经是numpy数组，则直接使用
            image_embedding = embedding
        # 创建一个内存中的二进制流
        buffer = io.BytesIO()
        # 将numpy数组保存到二进制流中，格式为.npy
        np.save(buffer, image_embedding)
        # 将文件指针重置到流的开头
        buffer.seek(0)
        # 返回包含.npy数据的响应，并添加必要的响应头
        return Response(
            content=buffer.getvalue(),
            media_type="application/octet-stream",
            headers={
                "X-Image-ID": image_id,
                "X-Model-Type": config.MODEL_TYPE,
                "X-Embedding-Shape": ",".join(map(str, embedding.shape)),
                "Content-Disposition": f"attachment; filename=embedding_{image_id}.npy"
            }
        )
       
    except HTTPException as e:
        # 重新抛出HTTP异常
        raise e
    except Exception as e:
        logger.error(f"生成embedding失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 新增/embedding/upload接口，专门用于文件上传
@app.post("/embedding/upload", response_model=EmbeddingResponse)
async def generate_embedding_upload(
    file: UploadFile = File(..., description="要处理的图像文件")
):
    """上传图像文件，生成embedding（专用接口）"""
    try:
        # 检查文件类型
        if not file.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="上传的文件必须是图像")
        
        # 读取文件内容
        image_bytes = await file.read()
        
        # 生成唯一的图像ID
        image_id = str(uuid.uuid4())
        
        logger.info(f"接收到图像上传请求，文件: {file.filename}, 大小: {len(image_bytes)} bytes, ID: {image_id}")
        
        # 生成embedding
        embedding = sam_model.generate_embedding(image_bytes, image_id)
        
        # 将numpy数组转换为列表
        embedding_list = embedding.tolist()
        
        logger.info(f"成功生成embedding，形状: {embedding.shape}, ID: {image_id}")
        
        return EmbeddingResponse(
            embedding=embedding_list,
            shape=list(embedding.shape),
            image_id=image_id,
            model_type=config.MODEL_TYPE
        )
    except Exception as e:
        logger.error(f"生成embedding失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/cache")
def clear_cache():
    """清除缓存"""
    try:
        sam_model.clear_cache()
        return {"status": "ok", "message": "缓存已清除"}
    except Exception as e:
        logger.error(f"清除缓存失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/stats")
def get_stats():
    """获取服务统计信息"""
    try:
        stats = {
            "model_loaded": sam_model.model_loaded,
            "device": str(sam_model.device),
            "cache_enabled": config.ENABLE_CACHE,
            "cache_size": len(sam_model.cache) if config.ENABLE_CACHE else 0,
            "max_cache_size": config.CACHE_SIZE if config.ENABLE_CACHE else 0,
            "model_type": config.MODEL_TYPE
        }
        return stats
    except Exception as e:
        logger.error(f"获取统计信息失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 如果直接运行该文件，启动服务
if __name__ == "__main__":
    import uvicorn
    
    logger.info(f"启动SAM Embedding服务，监听地址: {config.HOST}:{config.PORT}, 设备: {config.DEVICE}")
    
    # 预加载模型
    sam_model.load_model()
    
    # 启动服务
    uvicorn.run(
        "api:app",
        host=config.HOST,
        port=config.PORT,
        reload=False,
        workers=1,
        log_level=config.LOG_LEVEL.lower()
    )